# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

from dataclasses import dataclass
from typing import Any

import torch

from forge.data_models.prompt import Prompt


@dataclass
class Completion:
    """A model-generated completion for a given prompt."""

    # The original prompt.
    prompt: Prompt

    # the decoded text returned by the model
    text: str

    # the encoded text (token ids) that were fed into the model
    prompt_ids: torch.Tensor

    # the encoded text (token ids) that were generated by the model
    token_ids: torch.Tensor

    # the log probabilities of the target tokens
    logprobs: torch.Tensor | None = None

    # the reason for stopping the generation
    stop_reason: str | None = None

    # the version identifier of the model when the generation was performed
    generator_version: int | None = None

    # extra information that might be useful for debugging
    metadata: dict[str, Any] | None = None
